import os
import shutil
from pathlib import Path
from typing import Optional

SOURCE_DIR = "/ssd/data/swp/RAVDESS/Audio_Speech_Actors_01-24"
DEST_BASE = "/hdd/data/swp/DH/RAVDESS/audio"


EMOTION_TO_CATEGORY = {
    "01": "neutral",   # neutral
    "02": "neutral",   # calm -> neutral
    "03": "positive",  # happy -> positive
    "08": "positive",  # surprised -> positive
    "04": "negative",  # sad -> negative
    "05": "negative",  # angry -> negative
    "06": "negative",  # fearful -> negative
    "07": "negative",  # disgust -> negative
}


def ensure_dest_dirs(base_dir: str) -> None:
    for sub in {"neutral", "positive", "negative"}:
        Path(base_dir, sub).mkdir(parents=True, exist_ok=True)


def extract_emotion_code(filename: str) -> Optional[str]:
    name = os.path.basename(filename)
    if not name.lower().endswith(".wav"):
        return None
    stem = os.path.splitext(name)[0]
    parts = stem.split("-")
    if len(parts) < 3:
        return None
    return parts[2]


def categorize_emotion(emotion_code: str) -> Optional[str]:
    return EMOTION_TO_CATEGORY.get(emotion_code)


def copy_wav_files_by_emotion(source_dir: str, dest_base: str) -> None:
    ensure_dest_dirs(dest_base)

    src_path = Path(source_dir)
    if not src_path.exists():
        raise FileNotFoundError(f"Source directory not found: {source_dir}")

    total = 0
    copied = 0
    skipped = 0

    for wav_path in src_path.rglob("*.wav"):
        total += 1
        emotion_code = extract_emotion_code(wav_path.name)
        if not emotion_code:
            skipped += 1
            print(f"[SKIP] Cannot parse emotion from filename: {wav_path}")
            continue

        category = categorize_emotion(emotion_code)
        if not category:
            skipped += 1
            print(f"[SKIP] Unknown emotion code '{emotion_code}' for file: {wav_path}")
            continue

        dest_dir = Path(dest_base, category)
        dest_file = dest_dir / wav_path.name

        try:
            shutil.copy2(wav_path, dest_file)
            copied += 1
            print(f"[OK] {wav_path} -> {dest_file}")
        except Exception as e:
            skipped += 1
            print(f"[ERR] Failed to copy {wav_path} -> {dest_file}: {e}")

    print(
        f"\nDone. Total: {total}, Copied: {copied}, Skipped: {skipped}. "
        f"Dest base: {dest_base}"
    )


if __name__ == "__main__":
    copy_wav_files_by_emotion(SOURCE_DIR, DEST_BASE)


